package com.heima.common.llm;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import lombok.Getter;
import lombok.Setter;
import org.apache.http.HttpEntity;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.StringEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.stereotype.Component;

import javax.imageio.ImageIO;
import java.awt.image.BufferedImage;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.HashMap;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Getter
@Setter
@Component
@ConfigurationProperties(prefix = "ocr")
public class OCR {

    private String apiKey;
    private String apiUrl;
    // 注意：glm-4.1v-thinking-flash擅长场景分析，OCR建议优先用glm-4.5v
    private String model;
    private boolean doSample = false;
    private boolean stream = false;
    private float temperature = 0.0f; // 温度0确保OCR结果精准无推测
    private float topP = 0.95f;

    /**
     * 核心工具方法：将BufferedImage转换为Base64字符串
     */
    public static String convertImageToBase64Str(BufferedImage image) {
        if (image == null) {
            throw new IllegalArgumentException("BufferedImage cannot be null");
        }

        ByteArrayOutputStream baos = null;
        try {
            baos = new ByteArrayOutputStream();

            // 尝试常用图片格式进行编码
            String[] formats = {"png", "jpg", "jpeg", "gif", "bmp"};
            for (String format : formats) {
                if (ImageIO.write(image, format, baos)) {
                    // 构建完整的data URI格式，包含MIME类型
                    byte[] bytes = baos.toByteArray();
                    String base64 = Base64.getEncoder().encodeToString(bytes);
                    return "data:image/" + format + ";base64," + base64;
                }
            }

            throw new IOException("Unsupported image type or could not write image");
        } catch (Exception e) {
            e.printStackTrace();
            return null;
        } finally {
            if (baos != null) {
                try {
                    baos.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
    }


    /**
     * 识别单张图片中的文字（输入BufferedImage对象）
     */
    public Map<String, String> doOCR(BufferedImage image) {
        // 转换图片为Base64
        String imageBase64 = convertImageToBase64Str(image);
        if (imageBase64 == null) {
            Map<String, String> errorMap = new HashMap<>();
            errorMap.put("status", "error");
            errorMap.put("content", "图片转换为Base64失败");
            return errorMap;
        }

        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
            HttpPost httpPost = new HttpPost(apiUrl);
            httpPost.setHeader("Content-Type", "application/json");
            httpPost.setHeader("Authorization", "Bearer " + apiKey);
            httpPost.setHeader("Connection", "close"); // 避免长连接超时

            JSONObject requestBody = new JSONObject();
            requestBody.put("model", model);
            requestBody.put("do_sample", doSample);
            requestBody.put("stream", stream);
            requestBody.put("temperature", temperature);
            requestBody.put("top_p", topP);

            // 响应格式：强制纯文本，避免模型返回特殊标签
            JSONObject responseFormat = new JSONObject();
            responseFormat.put("type", "text");
            requestBody.put("response_format", responseFormat);

            // 构建消息体：图片（带正确Base64） + 提示词
            JSONArray messages = new JSONArray();
            JSONObject message = new JSONObject();
            message.put("role", "user");
            JSONArray contentArray = new JSONArray();

            // 处理图片
            JSONObject imageContent = new JSONObject();
            imageContent.put("type", "image_url");

            JSONObject imageUrlObj = new JSONObject();
            imageUrlObj.put("url", imageBase64); // 传入完整格式的图片地址
            imageContent.put("image_url", imageUrlObj);

            contentArray.add(imageContent);

            // 处理提示词：明确绑定图片，避免模型偏离
            JSONObject textContent = new JSONObject();
            textContent.put("type", "text");
            textContent.put("text", "请识别图片中的所有文字，包括标点符号，保持原文格式和顺序，直接返回识别结果，不需要额外解释。");
            contentArray.add(textContent);

            message.put("content", contentArray);
            messages.add(message);
            requestBody.put("messages", messages);


            // 设置请求体（UTF-8编码，避免中文乱码）
            httpPost.setEntity(new StringEntity(
                    requestBody.toString(),
                    StandardCharsets.UTF_8.name()
            ));

            // 发送请求并处理响应
            try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
                HttpEntity entity = response.getEntity();
                if (entity != null) {
                    String responseBody = EntityUtils.toString(entity, StandardCharsets.UTF_8);

                    if (response.getStatusLine().getStatusCode() == 200) {
                        JSONObject result = JSON.parseObject(responseBody);
                        Map<String, String> resultMap = new HashMap<>();

                        JSONArray choices = result.getJSONArray("choices");
                        if (choices != null && !choices.isEmpty()) {
                            JSONObject firstChoice = choices.getJSONObject(0);
                            JSONObject messageResult = firstChoice.getJSONObject("message");
                            String content = messageResult.getString("content");

                            // 后处理：移除模型可能残留的特殊标签
                            if (content != null) {
                                content = content.replaceAll("<\\|.*?\\|>", "")
                                        .replaceAll("\\n+", "\n")
                                        .trim();
                            }

                            resultMap.put("status", "success");
                            resultMap.put("content", content);
                            return resultMap;
                        }

                        resultMap.put("status", "empty");
                        resultMap.put("content", "");
                        return resultMap;
                    } else {
                        // API错误：返回详细信息方便调试（如密钥无效、额度不足）
                        Map<String, String> errorMap = new HashMap<>();
                        errorMap.put("status", "error");
                        errorMap.put("content", "API请求失败，状态码: " + response.getStatusLine().getStatusCode());
                        errorMap.put("details", responseBody);
                        return errorMap;
                    }
                }
            }
        } catch (IOException e) {
            Map<String, String> errorMap = new HashMap<>();
            errorMap.put("status", "exception");
            errorMap.put("content", "网络请求异常：" + e.getMessage());
            return errorMap;
        }

        // 极端情况：返回空字符串+异常状态
        Map<String, String> defaultErrorMap = new HashMap<>();
        defaultErrorMap.put("status", "exception");
        defaultErrorMap.put("content", "");
        return defaultErrorMap;
    }

    /**
     * 批量识别多张图片中的文字（输入BufferedImage列表）
     */
    public Map<String, String> doOCRBatch(List<BufferedImage> images) {
        List<String> imageBase64List = new ArrayList<>();
        for (BufferedImage image : images) {
            String base64 = convertImageToBase64Str(image);
            if (base64 != null) {
                imageBase64List.add(base64);
            }
        }

        if (imageBase64List.isEmpty()) {
            Map<String, String> errorMap = new HashMap<>();
            errorMap.put("status", "error");
            errorMap.put("content", "所有图片转换为Base64均失败");
            return errorMap;
        }

        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
            HttpPost httpPost = new HttpPost(apiUrl);
            httpPost.setHeader("Content-Type", "application/json");
            httpPost.setHeader("Authorization", "Bearer " + apiKey);
            httpPost.setHeader("Connection", "close");

            JSONObject requestBody = new JSONObject();
            requestBody.put("model", model);
            requestBody.put("do_sample", doSample);
            requestBody.put("stream", stream);
            requestBody.put("temperature", temperature);
            requestBody.put("top_p", topP);

            JSONObject responseFormat = new JSONObject();
            responseFormat.put("type", "text");
            requestBody.put("response_format", responseFormat);

            JSONArray messages = new JSONArray();
            JSONObject message = new JSONObject();
            message.put("role", "user");
            JSONArray contentArray = new JSONArray();

            // 添加所有图片
            for (String imageBase64 : imageBase64List) {
                JSONObject imageContent = new JSONObject();
                imageContent.put("type", "image_url");

                JSONObject imageUrlObj = new JSONObject();
                imageUrlObj.put("url", imageBase64);
                imageContent.put("image_url", imageUrlObj);

                contentArray.add(imageContent);
            }

            // 添加提示词
            JSONObject textContent = new JSONObject();
            textContent.put("type", "text");
            textContent.put("text", "请识别所有图片中的文字，包括标点符号，保持原文格式和顺序，按图片顺序依次返回识别结果，不需要额外解释。");
            contentArray.add(textContent);

            message.put("content", contentArray);
            messages.add(message);
            requestBody.put("messages", messages);

            httpPost.setEntity(new StringEntity(
                    requestBody.toString(),
                    StandardCharsets.UTF_8.name()
            ));

            try (CloseableHttpResponse response = httpClient.execute(httpPost)) {
                HttpEntity entity = response.getEntity();
                if (entity != null) {
                    String responseBody = EntityUtils.toString(entity, StandardCharsets.UTF_8);

                    if (response.getStatusLine().getStatusCode() == 200) {
                        JSONObject result = JSON.parseObject(responseBody);
                        Map<String, String> resultMap = new HashMap<>();

                        JSONArray choices = result.getJSONArray("choices");
                        if (choices != null && !choices.isEmpty()) {
                            JSONObject firstChoice = choices.getJSONObject(0);
                            JSONObject messageResult = firstChoice.getJSONObject("message");
                            String content = messageResult.getString("content");

                            if (content != null) {
                                content = content.replaceAll("<\\|.*?\\|>", "")
                                        .replaceAll("\\n+", "\n")
                                        .trim();
                            }

                            resultMap.put("status", "success");
                            resultMap.put("content", content);
                            return resultMap;
                        }

                        resultMap.put("status", "empty");
                        resultMap.put("content", "");
                        return resultMap;
                    } else {
                        Map<String, String> errorMap = new HashMap<>();
                        errorMap.put("status", "error");
                        errorMap.put("content", "API请求失败，状态码: " + response.getStatusLine().getStatusCode());
                        errorMap.put("details", responseBody);
                        return errorMap;
                    }
                }
            }
        } catch (IOException e) {
            Map<String, String> errorMap = new HashMap<>();
            errorMap.put("status", "exception");
            errorMap.put("content", "网络请求异常：" + e.getMessage());
            return errorMap;
        }

        Map<String, String> defaultErrorMap = new HashMap<>();
        defaultErrorMap.put("status", "exception");
        defaultErrorMap.put("content", "");
        return defaultErrorMap;
    }
}
